from functools import lru_cache
from pathlib import Path
from typing import Optional, ClassVar

import attrs
import pandas as pd
import yaml
from typing_extensions import Self

from .iterablenamespace import FrozenList


# OKLAHOMA_FIPS: str = "40"
_COUNTY_INFO_YAML_PATH = Path(fr"{__file__}\..\..\..\counties.yml")


class CountyInfoDumper(yaml.Dumper):
    """
    Class used in the process of generating ``counties.yml``.

    .. seealso::

       :meth:`CountyInfo._yaml_from_feature_class`
    """
    def represent_data(self, data):
        if isinstance(data, CountyInfo):
            # noinspection PyTypeChecker
            data = attrs.asdict(data)
        return super().represent_data(data)


@attrs.frozen(eq=True, order=True)
class CountyInfo:
    """
    Provides some basic information about counties based on the data in
    ``counties.yml``. The class method :meth:`load` automatically reads in the
    data from that YAML file and returns a :class:`FrozenList` of
    ``CountyInfo`` instances, while the class method :meth:`load_df` provides
    that information as a ``pandas.DataFrame``.

    When sorted, instance of this class are arranged alphabetically by their
    :attr:`state`, then by :attr:`name`, then by :attr:`fips3`, then by
    :attr:`number`.
    """

    state: str = attrs.field(validator=attrs.validators.matches_re(r"^[A-Z ]+$"))
    """State in which the county lies (uppercase)."""

    name: str = attrs.field(validator=attrs.validators.matches_re(r"^[A-Z ]+$"))
    """Name of the county (uppercase)."""

    fips3: str = attrs.field(validator=attrs.validators.matches_re(r"^\d{3}$"))
    """Three-digit FIPS code for the county."""

    number: int | None = attrs.field(default=None)
    """State-assigned county number for the county."""

    _county_equivalent_name_default: ClassVar[str] = "COUNTY"
    """Term used by default to refer to county-equivalent jurisdictions in
    computing :attr:`full_name`. Overridden for individual states in
    :attr:`_county_equivalent_names`."""

    _county_equivalent_names: ClassVar[dict[str, str]] = {}
    """Mapping of state name (:attr:`state`) to the term used locally for county-equivalent
    jurisdictions, e.g., ``{"LOUISIANA": "PARISH"}``. For states not included
    in this dict, the string in :attr:`_county_equivalent_name_default` will be
    used by :attr:`full_name`."""

    @classmethod
    @lru_cache(1)
    def load(cls) -> FrozenList[Self]:
        """Loads and returns all ``CountyInfo`` objects represented in
        ``counties.yml``."""
        with open(_COUNTY_INFO_YAML_PATH, "r") as yaml_file:
            yaml_data = yaml.safe_load(yaml_file)
        return FrozenList([cls(**field_data) for field_data in yaml_data])

    @classmethod
    def load_df(cls) -> pd.DataFrame:
        """
        Loads and returns all ``CountyInfo`` objects represented in
        ``counties.yml`` as a ``pandas.DataFrame``. The resulting data frame
        will have the following columns and ``dtype``\ s:

        * ``state`` - ``string[python]``
        * ``name`` - ``string[python]``
        * ``fips3`` - ``string[python]``
        * ``number`` - ``Int64``
        """
        # noinspection PyTypeChecker
        return pd.DataFrame(
            map(attrs.asdict, cls.load())
        ).astype({
            "state": pd.StringDtype(),
            "name": pd.StringDtype(),
            "fips3": pd.StringDtype(),
            "number": pd.Int64Dtype()
        })

    @classmethod
    def get_county(cls, name: Optional[str] = None, state: Optional[str] = None, *, fips3: Optional[str | int] = None, number: Optional[int] = None) -> Self:
        """
        Returns the ``CountyInfo`` object matching the supplied criteria.

        Exactly one of *name*, *fips3*, or *number* is required. If multiple
        counties match based on the given criterion of those three options, the
        *state* argument will be used to filter further. If a *state* argument
        is needed but absent, or if the number of candidates is not exactly
        one, a ``LookupError`` will be raised.

        In the case of *name*, a space followed by the jurisdiction type name,
        e.g., "COUNTY", is automatically removed. If *state* is provided,
        :attr:`_county_equivalent_names` will be used to determine the term to
        strip. If *state* is not provided,
        :attr:`_county_equivalent_name_default` will be used.

        :param name: Name of the county to search for
        :type name: Optional[str]
        :param state: Name of the state containing the desired county
        :type state: Optional[str]
        :param fips3: Three-digit FIPS code for the desired county
        :type fips3: Optional[str | int]
        :param number: The state-assigned county number of the desired county
        :type number: Optional[int]
        :return: The county matching the supplied criteria
        :rtype: CountyInfo
        """
        if len([*filter(bool, [name, fips3, number])]) != 1:
            raise TypeError("Exactly one of 'name', 'fips3', or 'number' must be provided.")
        values = cls.load()

        # Initial filtering
        if name:
            name = name.upper().rstrip(f" {cls._county_equivalent_names.get(state, cls._county_equivalent_name_default)}")
            candidates = [x for x in values if x.name == name]
        elif fips3:
            candidates = [x for x in values if x.fips3 == fips3]
        elif number:
            candidates = [x for x in values if x.number == number]
        else:
            candidates = []

        # Refine by state if needed
        if len(candidates) > 1:
            if not state:
                raise LookupError("Multiple counties matched the given criteria, but no argument was provided for 'state'.")
            candidates = [x for x in candidates if x.state == state.upper()]

        # Return result
        if len(candidates) == 1:
            return candidates[0]
        else:
            raise LookupError(f"Unexpectedly matched {len(candidates)} object(s).")


    @classmethod
    def _yaml_from_feature_class(cls, input_fc: Path, output_yaml: Path, name_field: str, fips_field: str, number_field: Optional[str] = None, state_field: Optional[str] = None, default_state: Optional[str] = None) -> Path:
        """
        Generates ``counties.yml`` based on a feature class containing the
        relevant data.

        .. warning::
           This method will, without confirmation, overwrite the file located
           at *output_yaml*, should one already exist.

        :param input_fc: Feature class containing relevant data
        :type input_fc: pathlib.Path
        :param output_yaml: Path to which the generated YAML will be written
        :type output_yaml: pathlib.Path
        :param name_field: Name of the county name field in *input_fc*
        :type name_field: str
        :param fips_field: Name of the three-digit FIPS code field in
            *input_fc*
        :type fips_field: str
        :param number_field: Name of the state-assigned county number field in
            *input_fc*
        :type number_field: Optional[str]
        :param state_field: Name of the state  field in *input_fc*
        :type state_field: Optional[str]
        :param default_state: State to use for all records generated in lieu of
            a *state_field*
        :type default_state: Optional[str]
        :return: Path to the generated YAML file
        :rtype: pathlib.Path
        """
        # noinspection PyUnresolvedReferences
        from arcgis.features import GeoSeriesAccessor, GeoAccessor

        if state_field and not default_state:
            fields = [name_field, fips_field, number_field, state_field]
        elif default_state and not state_field:
            fields = [name_field, fips_field, number_field]
        else:
            raise TypeError("Exactly one of 'state_field' or 'default_state' is required.")
        df: pd.DataFrame = pd.DataFrame.spatial.from_featureclass(str(input_fc), fields=fields)
        df["$name"] = df[name_field].str.upper()
        df["$state"] = df[state_field].str.upper() if state_field else default_state
        df["$fips3"] = df[fips_field].apply(lambda fips_code: f"{int(fips_code):03d}")
        df["$number"] = df[number_field].apply(lambda co_num: int(co_num) if co_num is not None else pd.NA) if number_field else pd.NA
        county_infos: list[cls] = df.apply(lambda row: cls( row["$state"], row["$name"], row["$fips3"], row["$number"]), axis=1).to_list()
        county_infos.sort()

        with open(output_yaml, "w") as f:
            yaml.dump(county_infos, f, CountyInfoDumper, sort_keys=False)
        return output_yaml.absolute()

    @property
    def full_name(self) -> str:
        return f"{self.name} {self._county_equivalent_names.get(self.state, self._county_equivalent_name_default)}"


if __name__ == "__main__":
    import argparse
    _parser = argparse.ArgumentParser()
    _parser.add_argument("input_feature_class")
    _parser.add_argument("output_yaml")
    _state_group = _parser.add_mutually_exclusive_group(required=True)
    _state_group.add_argument("-s", "--state-field", type=str)
    _state_group.add_argument("-S", "--state-default", type=str)
    _parser.add_argument("-c", "--county-field", type=str)
    _parser.add_argument("-f", "--fips-field", type=str)
    _parser.add_argument("-n", "--number-field", type=str)
    _args = _parser.parse_args()
    # noinspection PyProtectedMember
    out_path = CountyInfo._yaml_from_feature_class(Path(_args.input_feature_class), Path(_args.output_yaml), _args.county_field, _args.fips_field, _args.number_field, _args.state_field, _args.state_default)
    print(f"Wrote output to {out_path}.")

__all__ = ["CountyInfo", "CountyInfoDumper"]